import numpy as np
import pytest
from astropy.constants import G, M_earth, R_earth

import ballistic  # импортируем модуль целиком, чтобы подменять константу

# Порог сходимости результатов 1%
MAX_RELATIVE_ERROR = 0.01

# Подмена значения в модуле ballistic
ballistic.SAT_DRY_MASS = 1

# Создание трехмерного единичного вектора
direction = np.array([1, 1, 1])
direction = direction / np.linalg.norm(direction)

# Test data
test_data = [
    pytest.param(
        (
            direction * R_earth.value,  # на поверхности Земли
            (0, 0, 0),  # тяга отключена
            0,  # масса топлива, влияет только на общую массу КА
            (0, 0, 0),  # скорость нулевая, не учитываем аэродинамическое сопротивление
        ),
        -direction * (G * M_earth / R_earth**2).value,
        id="Near the planet without thrust",
    ),
    pytest.param(
        (
            direction * 1e30,  # на сильном удалении от Земли, на "бесконечности"
            (0, 0, 0),  # тяга отключена
            0,  # масса топлива, влияет только на общую массу КА
            (0, 0, 0),  # скорость нулевая, не учитываем аэродинамическое сопротивление
        ),
        [0, 0, 0],
        id="Infinity without thrust",
    ),
    pytest.param(
        (
            direction * 1e30,  # на сильном удалении от Земли, на "бесконечности"
            (1, 1, 1),  # двигатель работает
            0,  # масса топлива, влияет только на общую массу КА
            (0, 0, 0),  # скорость нулевая, не учитываем аэродинамическое сопротивление
        ),
        [1, 1, 1],
        id="Infinity with thrust",
    ),
]


@pytest.mark.parametrize("sat_acc_func_parameters, expected_result", test_data)
def test_satellite_acceleration(sat_acc_func_parameters, expected_result):
    assert np.allclose(
        ballistic.satellite_acceleration(*sat_acc_func_parameters),
        expected_result,
        rtol=MAX_RELATIVE_ERROR,
        equal_nan=True,
    )
